# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags

import datetime
import os
import logging

import numpy
import pandas

import coin.strategy.mm.base.test_feed_arb_strategy as tfas

from concurrent.futures import ProcessPoolExecutor


def getcompkey(row):
  keys = [
      'bid0p', 'bid0q', 'ask0p', 'ask0q', 'bid1p', 'bid1q', 'ask1p', 'ask1q', 'tradep', 'tradeq'
  ]
  return tuple(row[keys])


def rowsame(rowa, rowb):
  return getcompkey(rowa) == getcompkey(rowb)


def lcs(X, Y, xkeys, ykeys, same):
  m = len(X)
  n = len(Y)
  L = [[0 for x in range(n + 1)] for x in range(m + 1)]

  # Following steps build L[m+1][n+1] in bottom up fashion. Note
  # that L[i][j] contains length of LCS of X[0..i-1] and Y[0..j-1]
  for i in range(m + 1):
    for j in range(n + 1):
      if i == 0 or j == 0:
        L[i][j] = 0
      elif xkeys[i - 1] == ykeys[j - 1]:
        L[i][j] = L[i - 1][j - 1] + 1
      else:
        L[i][j] = max(L[i - 1][j], L[i][j - 1])

  # Following code is used to print LCS
  index = L[m][n]

  # Create a character array to store the lcs string
  lcs = [""] * (index + 1)
  lcs[index] = ""

  # Start from the right-most-bottom-most corner and
  # one by one store characters in lcs[]
  i = m
  j = n
  while i > 0 and j > 0:

    # If current character in X[] and Y are same, then
    # current character is part of LCS
    if same(X[i - 1], Y[j - 1]):
      lcs[index - 1] = (X[i - 1], Y[j - 1])
      i -= 1
      j -= 1
      index -= 1

    # If not same, then find the larger of two and
    # go in the direction of larger value
    elif L[i - 1][j] > L[i][j - 1]:
      i -= 1
    else:
      j -= 1
  return lcs


def get_paramname(params):
  return "%s_%s_sw%s" % (
      params['sub_req'], params['num_feed_worker'], params['switch_slow_agent_frequency_sec'])


def compare_params_list(params_list):
  params_df_dict = {}
  for params in params_list:
    print(params['out_file'])
    params_df_dict[str(params)] = pandas.read_csv(params['out_file'])

  comp_param = params_list[-1]
  for base_param in params_list[:-1]:
    lcs_result = lcs([row for _, row in params_df_dict[str(base_param)].iterrows()],
                     [row for _, row in params_df_dict[str(comp_param)].iterrows()],
                     [getcompkey(row) for _, row in params_df_dict[str(base_param)].iterrows()],
                     [getcompkey(row) for _, row in params_df_dict[str(comp_param)].iterrows()],
                     rowsame)
    lcs_result = lcs_result[5:-5]
    tds = [lcsrow[0].timestamp - lcsrow[1].timestamp for lcsrow in lcs_result[:-1]]
    tds = [td * 1e-6 for td in tds]

    def concatrow(lcs_row):
      return pandas.concat([
          pandas.DataFrame(lcs_row[0]).T.add_suffix('1').T,
          pandas.DataFrame(lcs_row[1]).T.add_suffix('2').T
      ]).T

    concatted = [concatrow(lcs_row) for lcs_row in lcs_result]
    df = pandas.concat(concatted, axis=0)

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams['font.size'] = 6
    plt.rcParams['figure.figsize'] = 12, 8
    plt.rcParams['lines.markersize'] = 1

    plt.subplot(211)
    plt.plot(tds, 'g.', markersize=1)
    title = "%s - %s(unit=ms), ratio(>)=%s" % (get_paramname(base_param),
                                               get_paramname(comp_param),
                                               (numpy.array(tds) < 0).mean())
    plt.title(title)
    plt.hlines(y=0, color='y', linestyle=':', lw=0.5, xmin=0, xmax=len(tds))

    plt.subplot(212)
    plt.plot(numpy.array(range(len(tds))) / len(tds), sorted(tds), 'g-', linewidth=0.5)
    plt.title("sorted(%s)" % title)
    plt.hlines(y=0, color='y', linestyle=':', lw=0.5, xmin=0, xmax=1)

    plt.savefig(
        os.path.join("test_feedarb",
                     "%s.png"
                     % "_vs_".join([get_paramname(base_param), get_paramname(comp_param)])))
    plt.close()


def launch_params_list(params_list):
  futures = []
  with ProcessPoolExecutor(max_workers=len(params_list)) as executor:
    for params in params_list:
      futures.append(executor.submit(tfas.run_feed_arb_strategy, **params))
    for future in futures:
      future.result()


def get_params_list(sub_req, num_feed_workers, switch_slow_agent_frequency_secs, exit_after_sec):
  params_list = []
  for i, num_feed_worker in enumerate(num_feed_workers):
    if switch_slow_agent_frequency_secs is None:
      switch_slow_agent_frequency_sec = None
    else:
      switch_slow_agent_frequency_sec = int(switch_slow_agent_frequency_secs[i])
      if switch_slow_agent_frequency_sec == 0:
        switch_slow_agent_frequency_sec = None
    num_feed_worker = int(num_feed_worker)
    paramdict = dict(num_feed_worker=num_feed_worker,
                     sub_req=sub_req,
                     switch_slow_agent_frequency_sec=switch_slow_agent_frequency_sec,
                     exit_after=datetime.timedelta(seconds=exit_after_sec),
                     quiet=True)
    out_file = os.path.join("test_feedarb", "%s.csv" % get_paramname(paramdict))
    paramdict['out_file'] = out_file
    params_list.append(paramdict)
  return params_list


def main(_):
  assert flags.FLAGS.num_feed_workers is not None
  assert flags.FLAGS.sub_req is not None
  assert flags.FLAGS.exit_after_sec is not None

  if flags.FLAGS.switch_slow_agent_frequency_secs is None:
    switch_slow_agent_frequency_secs = None
  else:
    switch_slow_agent_frequency_secs = flags.FLAGS.switch_slow_agent_frequency_secs.split(",")

  os.makedirs("test_feedarb", exist_ok=True)
  params_list = get_params_list(flags.FLAGS.sub_req,
                                flags.FLAGS.num_feed_workers.split(","),
                                switch_slow_agent_frequency_secs,
                                flags.FLAGS.exit_after_sec)
  if not flags.FLAGS.analyze_only:
    launch_params_list(params_list)
  compare_params_list(params_list)


if __name__ == '__main__':
  flags.DEFINE_string('num_feed_workers', None, '')
  flags.DEFINE_string('switch_slow_agent_frequency_secs', None, '')
  flags.DEFINE_string('sub_req', None, '')
  flags.DEFINE_integer('exit_after_sec', 30, '')
  flags.DEFINE_bool('analyze_only', False, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
